import gym
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from gym.utils import seeding
import random

import configparser
import tasks
from nodes import Node_U, Node_W1,Node_W2
import queue
import visdom


# matplotlib.use("Agg")  # 控制绘图不显示，必须在import matplotlib.pyplot as plt前运行
config = configparser.ConfigParser()
config.read('env_config.ini')
new_obs_dim = int(config.get('space', 'new_obs_dim'))
new_obs_dim = 14
new_action_dim = int(config.get('space', 'new_action_dim'))
new_action_dim = 3
class NetworkEnv(gym.Env):
    """

    """

    metadata = {"render.modes": ["human"]}

    '''
        初始化参数设置
    '''
    def __init__(
            self,
            args,                     # 其他参数
    ):
        # 重写观测空间维度
        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(new_obs_dim,))

        # 重写动作空间维度
        self.action_space = gym.spaces.Discrete(new_action_dim)

        # 初始化配置参数设置
        self.computing_node = args.NumOfWorkers
        self.consumer_node = args.NumOfUsers
        self.done_step = args.totalStep  #一个回合分配的时间步
        self.step_num = 0 #当前时间步
        self.seed = args.RamSeed
        random.seed(self.seed)
        # self.fig, self.ax = plt.subplots(2, 3, figsize=(14, 12))
        self.w1_node_num = 2
        ## 画图相关参数配置
        # for i in range(2):
        #     for j in range(3):
        #         self.ax[i][j].set_ylim([0, 100])
        #         self.ax[i][j].set_xlabel('Step')
        #         self.ax[i][j].set_ylabel('Value')
        #         self.ax[i][j].set_title(f'Worker {i}/{j} State Changes')
        #         self.ax[i][j].plot([], [], '-b')
        #         self.ax[i][j].plot([], [], '-r')
        #         self.ax[i][j].legend(['compute', 'band'])
        #
        # plt.ion()  # 开启交互模式
        self.vis = visdom.Visdom(env=args.visenv)

        self.args = args

        self.computing_ca = [15, 20, 30]  # 算力可选列表
        self.bandwidth_ca = [9, 20, 30, 50]  # 算力带宽列表

        self.computing_re = [10]  # index=10时，用户初始化
        self.bandwidth_re = [10]  # index=10，用户初始化
        
        # 初始化固定参数设置
        self.worker1_list = []       # 工作节点类列表
        self.worker2_list = []  # 工作子节点类列表
        self.user_list = []
        self.task_queue = queue.Queue()

        self.cl = [[[0] for _ in range(3)] for _ in range(2)]             # 存储状态变量 计算负载 的变化情况
        self.bl = [[[0] for _ in range(3)] for _ in range(2)]                       # 存储状态变量  带宽负载 的变化情况

        for index in range(self.w1_node_num,self.computing_node):  #初始化二级节点 2-5
            self.worker2_list.append(
                Node_W2(
                    self.computing_ca[index % len(self.computing_ca)], #30,15,20,30
                    self.bandwidth_ca[index % len(self.bandwidth_ca)],#  30,50,9,20
                    index,
                    self.args.NodeStrategyCo,
                    self.args.NodeStrategyTr
                )
            )

        self.task_queue_w1_list = [] #一级节点的任务队列
        for index in range(self.w1_node_num):
            self.task_queue_w1_list.append(queue.Queue())


        for index in range(self.w1_node_num):  #初始化一级节点,并为其分配子节点 0,1
            self.worker1_list.append(
                Node_W1(
                    self.computing_ca[index % len(self.computing_ca)], # 15,20
                    self.bandwidth_ca[index % len(self.bandwidth_ca)], # 9,20
                    index,
                    self.args.NodeStrategyCo,
                    self.args.NodeStrategyTr,[self.worker2_list[index*2],self.worker2_list[index*2+1]]
                )
            )
        self.work1_group = []  #二维数组，0维表示一级节点组号，第1维表示包括一级节点及子节点的节点类
        self.work1_group_id = []
        for index in range(self.w1_node_num):
            work1 = self.worker1_list[index]
            work2_list = work1.worker2
            assert isinstance(work2_list, list)
            work_group = [work1] + work2_list
            self.work1_group_id.append([work.nodeID for work in work_group])
            self.work1_group.append(work_group)

        print('分组id',self.work1_group_id)
        for index in range(self.consumer_node):
            self.user_list.append(
                Node_U(
                    index,index % 10,
                    self.computing_re[0],self.bandwidth_re[0]
                )
            )   # 通过 index % 10+1 选择不同的任务，范围 10-100 或自设值
    def info(self):
        for j,g in enumerate(self.work1_group):
            print(f'第{j}组')
            for i,w in enumerate(g):
                if i==0:
                    print(f'主节点{w.nodeID}，计算能力{w.computing_ca}，传输能力{w.bandwidth_ca},共处理任务{w.task_num},计算传输比为：{w.computing_ca/w.bandwidth_ca}->{w.task_process_num/(w.task_transmit_num+1e-6)}')
                    continue
                print(f'从节点{w.nodeID}，计算能力{w.computing_ca}，传输能力{w.bandwidth_ca},共处理任务{w.task_num},计算传输比为：{w.computing_ca/w.bandwidth_ca}->{w.task_process_num/(w.task_transmit_num+1e-6)}')
            print('\n')




    def step1(self, action):  #一级分配,决定分配到哪个工作组
        # print(self.task_queue.qsize())
        task = self.task_queue.get()
        # print(task.taskID)
        self.task_queue_w1_list[action].put(task)
        state_next = self.get_state()  #两个一级状态
        state_next[2+4*action+2] +=task.computing_re  #增加分配任务的计算和传输量
        state_next[2 + 4 * action + 3] += task.bandwidth_re
        r = None # 待任务完成后计算奖励;
        done=0
        tid = task.taskID


        return state_next,r,done,tid

    def stepbyworker1(self,worker1_id,action):
        task = self.task_queue_w1_list[worker1_id].get()
        self.work1_group[worker1_id][action].allocate_task('C', task)  #
        state_next = self.get_group_state(worker1_id)  # 两个一级状态
        r = None  # 待任务完成后计算奖励;
        done = 0
        tid = task.taskID
        return state_next, r, done, tid




    def reset(self):
        # 初始化固定参数设置

        self.task_queue.queue.clear() #情况任务队列
        self.cl = [[[0] for _ in range(3)] for _ in range(2)]      # 存储状态变量 计算负载 的变化情况
        self.bl = [[[0] for _ in range(3)] for _ in range(2)]     # 存储状态变量  带宽负载 的变化情况
        self.step_num = 0  # 时间步归0

        for i in range(len(self.worker1_list)): #清空节点队列
            self.worker1_list[i].task_process_list = []
            self.worker1_list[i].task_transmit_list = []
            self.worker1_list[i].task_num=0
            self.worker1_list[i].task_transmit_num=0
            self.worker1_list[i].task_process_num=0

            self.task_queue_w1_list[i].queue.clear()

        for i in range(len(self.worker2_list)):  # 清空节点队列
            self.worker2_list[i].task_process_list = []
            self.worker2_list[i].task_transmit_list = []
            self.worker2_list[i].task_num=0
            self.worker2_list[i].task_transmit_num=0
            self.worker2_list[i].task_process_num=0


        state = self.get_state()
        return state




    def get_state(self):  #获取上层状态
        if self.task_queue.qsize()==0:
            state =[0,0]
        else:
            task = self.task_queue.queue[0]
            state = [task.computing_re,task.bandwidth_re]  #待分配状态
        for node in self.worker1_list:
            cc, bc, rc, rb = node.get_all_node_state() #获取全局子节点
            state.extend([cc, bc, rc, rb])

        # print(f'开始分配前，状态信息为：{state}')
        return state

    def get_group_state(self,work_groupid):  #获取下层状态
        if self.task_queue_w1_list[work_groupid].qsize()==0:
            state =[0,0]
        else:
            task = self.task_queue_w1_list[work_groupid].queue[0]
            state = [task.computing_re,task.bandwidth_re]  #待分配状态
        for node in self.work1_group[work_groupid]:
            cc, bc, rc, rb = node.get_node_state() #获取全局子节点
            state.extend([cc, bc, rc, rb])

        # print(f'开始分配前，状态信息为：{state}')
        return state



    def check_idle(self): #检测是否空闲
        allrc=0
        allrb=0
        if self.task_queue.qsize() == 0:
            for node in self.worker_list:
                cc, bc, rc, rb = node.get_node_state()
                allrb +=rb
                allrc +=rc
            if allrb==0 and allrc==0:
                return True
        return False


    # 将用户产生的任务放在 task_queue 内
    def add_task(self, task):
        self.task_queue.put(task)

    # 画图函数
    def render(self, ep_i,mode="human", close=False):
        # for i in range(2):
        #     for j in range(3):
        #
        #
        #         self.ax[i][j].set_xlim([0, len(self.cl[i][j]) if len(self.cl[i][j]) > 100 else 100])  #
        #         self.ax[i][j].set_ylim([0, 100 if len(self.cl[i][j])==0 else max(max(self.cl[i][j]),max(self.bl[i][j]))])
        #         self.ax[i][j].lines[0].set_xdata(range(len(self.cl[i][j])))
        #         self.ax[i][j].lines[1].set_xdata(range(len(self.bl[i][j])))
        #         self.ax[i][j].lines[0].set_ydata(self.cl[i][j])
        #         self.ax[i][j].lines[1].set_ydata(self.bl[i][j])
        #         self.ax[i][j].relim()
        #         self.ax[i][j].autoscale_view(True, True, True)
        # plt.draw()
        # if mode == 'human':
        #     plt.pause(0.1)
        # 暂停一小段时间，允许交互窗口更新
        num_rows = 2
        num_cols = 3
        subplot_titles =[]
        for i in range(len(self.work1_group)):
            for j in range(len(self.work1_group[0])):
                subplot_titles.append(f'回合{ep_i},Computer {self.work1_group[i][j].nodeID},计算能力{self.work1_group[i][j].computing_ca}，传输能力{self.work1_group[i][j].bandwidth_ca}')


        # Create subplots for each computer's load
        for i in range(num_rows * num_cols):
            row = i // num_cols
            col = i % num_cols
            cl = self.cl[row][col]
            bl = self.bl[row][col]
            x1 = np.arange(len(self.cl[row][col]))

            self.vis.line(
                X=x1,
                Y=np.column_stack((cl, bl)),
                opts=dict(
                    title=subplot_titles[i],
                    xlabel='时间片',
                    ylabel='负载',
                    legend=['计算负载', '传输负载']
                ),
                win=subplot_titles[i]  # Each subplot is given a unique window name
            )


    def _seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]




